#!/usr/bin/env python
# coding=utf-8

import joblib
import os
import kaldi_io
from sklearn.decomposition import PCA
import numpy as np
import argparse
import random

parser = argparse.ArgumentParser()
parser.add_argument("--scp-path", type=str, default="dvector.scp", help="scp file path")
parser.add_argument("--out-dir", type=str, default="", help="PCA out dir")
args = parser.parse_args()

first = True
for idx, (key,mat) in enumerate(kaldi_io.read_mat_scp(args.scp_path)):
	mat = np.mean(mat, 0)
	if first == True:
		dvectors = mat.reshape((1, -1))
	else:
		dvectors = np.vstack((dvectors, mat.reshape((1, -1))))
	print("{} \ {}".format(idx, dvectors.shape))
	first = False
	if idx == 1000:
		break

		
print(dvectors.shape)

# PCA
components = [32, 64, 128, 256, 300 , 400]
for c in components:
	print("train PCA {}".format(c))
	pca = PCA(n_components=c, svd_solver="full")
	pca.fit(dvectors)
	dvectors = dvectors[:100]
	print(dvectors.shape)
	data = pca.fit_transform(dvectors)
	input(data.shape)
	dst_path = os.path.join(args.out_dir, "pca_{}.model".format(c))
	joblib.dump(pca, dst_path)

